#!/usr/bin/python3

import json
import socket
import sys

class QMPClient:
    def __init__(self, socket_path):
        self.socket_path = socket_path
        self.sock = None

    def connect(self):
        self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.sock.connect(self.socket_path)
        self._initialize()

    def _initialize(self):
        self._read_response()
        self._send_command({"execute": "qmp_capabilities"})

    def _read_response(self):
        response = b""
        while True:
            chunk = self.sock.recv(4096)
            if not chunk:
                break
            response += chunk
            if b'\n' in chunk:
                break
        return json.loads(response.decode('ascii'))


    def _send_command(self, command):
        self.sock.sendall(json.dumps(command).encode('utf-8') + b'\n')
        return self._read_response()

    def add_virtio_dev_pci(self, id):
       # Add the virtio-mem-pci device
        device_command = {
               "execute": "device_add",
               "arguments": {
                   "driver": "virtio-blk-pci",
                   "id": 'vdb',
                   "bus": "pci.1",
                   "drive": id,
                   }
               }
        return self._send_command(device_command)

    def del_virtio_dev_pci(self, id):
       # Del the virtio-mem-pci device
        device_command = {
               "execute": "device_del",
               "arguments": {
                   "id": 'vdb',
                   }
               }
        return self._send_command(device_command)


    def add_virtio_disk_device(self, path):
        # Add the disk
        # "aio": "io_uring",
        memdev_command = {
                "execute": "blockdev-add",
                "arguments": {
                    "driver": "raw",
                    "file": {
                        "driver": "host_device",
                        "filename": path,
                        "aio":"io_uring",
                     },
                    "read-only": False,
                    "node-name": "disk-0",
                }
        }
        disk_response = self._send_command(memdev_command)
        print("disk add response:", disk_response)

    def query_virtio(self):
        command = {
                "execute":"query-blockdev"
                }
        return self._send_command(command)

    def close(self):
        if self.sock:
            self.sock.close()

if __name__ == "__main__":
    qmp_socket_path = "/tmp/qmp-sock"

    client = QMPClient(qmp_socket_path)
    client.connect()
    
    if sys.argv[1] == "qv":
        response = client.query_virtio()
        print(response)
    elif sys.argv[1] == "addd":
        disk = sys.argv[2]
        response = client.add_virtio_disk_device(disk)
        print(response)
    elif sys.argv[1] == "set":
        # Get memory
        response = client.add_virtio_dev_pci("disk-0")
        print("add blk response:", response)
    elif sys.argv[1] == "del":
        # Get memory
        response = client.del_virtio_dev_pci("disk-0")
        print("del blk response:", response)
    
    client.close()
